from vit_pytorch.model import ViT, VIT_B16_224